package net.xiaohanlin.smartmetric.monitor.jvm;

import java.lang.management.ManagementFactory;
import java.lang.management.ThreadInfo;
import java.lang.management.ThreadMXBean;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Set;

import com.codahale.metrics.Gauge;
import com.codahale.metrics.Metric;
import com.codahale.metrics.MetricSet;
import com.codahale.metrics.jvm.ThreadDeadlockDetector;

import net.xiaohanlin.smartmetric.metricname.JvmThreadMetricNameEnum;

public class ThreadStatesGaugeSet extends AbstractJvmMetric implements MetricSet {

	// do not compute stack traces.
	private final static int STACK_TRACE_DEPTH = 0;

	private final ThreadMXBean threads;
	private final ThreadDeadlockDetector deadlockDetector;

	/**
	 * Creates a new set of gauges using the default MXBeans.
	 */
	public ThreadStatesGaugeSet() {
		this(ManagementFactory.getThreadMXBean(), new ThreadDeadlockDetector());
	}

	/**
	 * Creates a new set of gauges using the given MXBean and detector.
	 *
	 * @param threads a thread MXBean
	 * @param deadlockDetector a deadlock detector
	 */
	public ThreadStatesGaugeSet(ThreadMXBean threads, ThreadDeadlockDetector deadlockDetector) {
		this.threads = threads;
		this.deadlockDetector = deadlockDetector;
	}

	@Override
	public Map<String, Metric> getMetrics() {
		final Map<String, Metric> gauges = new HashMap<String, Metric>();

		for (final Thread.State state : Thread.State.values()) {
			gauges.put(getMetricName(JvmThreadMetricNameEnum.valueOf("jvm_thread_state_" + state.toString().toLowerCase())), new Gauge<Object>() {

				@Override
				public Object getValue() {
					return getThreadCount(state);
				}
			});
		}

		gauges.put(getMetricName(JvmThreadMetricNameEnum.jvm_thread_state_new), new Gauge<Object>() {

			@Override
			public Object getValue() {
				return getThreadCount(Thread.State.NEW);
			}
		});

		gauges.put(getMetricName(JvmThreadMetricNameEnum.jvm_thread_count), new Gauge<Integer>() {

			@Override
			public Integer getValue() {
				return threads.getThreadCount();
			}
		});

		gauges.put(getMetricName(JvmThreadMetricNameEnum.jvm_thread_daemon_count), new Gauge<Integer>() {

			@Override
			public Integer getValue() {
				return threads.getDaemonThreadCount();
			}
		});

		gauges.put(getMetricName(JvmThreadMetricNameEnum.jvm_thread_deadlock_count), new Gauge<Integer>() {

			@Override
			public Integer getValue() {
				return deadlockDetector.getDeadlockedThreads().size();
			}
		});

		gauges.put(getMetricName(JvmThreadMetricNameEnum.jvm_thread_deadlocks), new Gauge<Set<String>>() {

			@Override
			public Set<String> getValue() {
				return deadlockDetector.getDeadlockedThreads();
			}
		});

		return Collections.unmodifiableMap(gauges);
	}

	private int getThreadCount(Thread.State state) {
		final ThreadInfo[] allThreads = getThreadInfo();
		int count = 0;
		for (ThreadInfo info : allThreads) {
			if (info != null && info.getThreadState() == state) {
				count++;
			}
		}
		return count;
	}

	ThreadInfo[] getThreadInfo() {
		return threads.getThreadInfo(threads.getAllThreadIds(), STACK_TRACE_DEPTH);
	}

}
